from random import *
from utils import *
from config import *

orientation_distance_templates = [
    "Mark the pixel coordinate for the location that is {distance} {direction} on the desktop.",
    "Specify the pixel coordinate for the point that is {distance} {direction} on the desktop.",
    "Indicate the pixel position at a distance of {distance} {direction} on the desktop.",
    "Point out the pixel coordinate of the location {distance} {direction} on the desktop.",
    "Where is the pixel location {distance} {direction} on the desktop?",
    "Could you mark the pixel coordinate for a spot {distance} {direction} on the desktop?",
    "Would you indicate the pixel position that is {distance} {direction} on the desktop?",
    "Identify the pixel coordinate of the point lying {distance} {direction} on the desktop.",
    "Please provide the pixel coordinate for the position {distance} {direction} on the desktop.",
    "Can you mark the pixel location that is {distance} {direction} on the desktop?",
    "Locate the pixel position for a point {distance} {direction} on the desktop.",
    "Reveal the pixel coordinate for a spot positioned {distance} {direction} on the desktop.",
    "What is the pixel coordinate for a location {distance} {direction} on the desktop?",
    "Determine the pixel position for the point {distance} {direction} on the desktop.",
    "Report the pixel coordinate for the position that is {distance} {direction} on the desktop.",
    "Show the pixel location of the point that lies {distance} {direction} on the desktop.",
    "Could you indicate the pixel coordinate at a location {distance} {direction} on the desktop?",
    "Please mark the pixel coordinate for the spot {distance} {direction} on the desktop.",
    "Would you please specify the pixel coordinate for a point {distance} {direction} on the desktop?",
    "Can you please reveal the pixel position for the location {distance} {direction} on the desktop?"
]
def gen_orientation_distance_qa(obj_dict, cam, table):
    qa = []
    for id, info in obj_dict.items():
        if not info['visible'] or not info['bottom_visible']:
            continue
        ref = choice(info['exclusive_ref'])
        desc = ref_to_desc(ref)
        dir = info['direction']
        dist_desc, dir_desc = None, None
        dirs = [
            ('left', (0, -1)),
            ('right', (0, 1)),
            ('back', (-1, 0)),
            ('front', (1, 0)),
        ]
        if not float_eq(abs(dir[0] + dir[1]), 1):
            times = 0
            while times < 8:
                times += 1
                dirid = randint(0, 3)
                distance = randint(5, 40) * 0.01
                pos = (info['bottom_points'][dirid][0] + dirs[dirid][1][0] * distance,
                       info['bottom_points'][dirid][1] + dirs[dirid][1][1] * distance)
                pix = check_table_pos(pos, cam, table)
                if pix == None:
                    continue
                dist_desc = f'{distance}m' if random() < 0.5 else ''
                dir_desc = f'on the {dirs[dirid][0]} of {desc}'
                break
        else:
            if float_eq(dir[1], -1):
                dirid = 0
            elif float_eq(dir[1], 1):
                dirid = 1
            elif float_eq(dir[0], -1):
                dirid = 2
            else:
                dirid = 3
            times = 0
            while times < 8:
                times += 1
                distance = randint(5, 40) * 0.01
                pos = (info['bottom_points'][dirid][0] + dirs[dirid][1][0] * distance,
                       info['bottom_points'][dirid][1] + dirs[dirid][1][1] * distance)
                pix = check_table_pos(pos, cam, table)
                if pix == None:
                    continue
                dist_desc = f'{distance}m' if random() < 0.5 else ''
                dir_desc = info['dir_desc'].replace('OBJECT', desc)
                break
        if pix != None:
            question = choice(orientation_distance_templates).format(
                distance = dist_desc,
                direction = dir_desc
            )
            if random() < 0.4:
                question = question + SUFFIX
            answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
            side_pix = world_to_camera_pixel(info["bottom_points"][dirid], cam)
            thinking = ref_to_thinking(obj_dict, info, ref, desc) + '\n' + f'''
The coordinates of the bottom point of the object on the {dirs[dirid][0]} side is {format_coordinatexy(side_pix[0], side_pix[1])}.
So the point {dist_desc} {dir_desc} is {format_coordinatexy(pix[0], pix[1])}.
            '''
            qa.append({
                'question': question,
                'answer': answer,
                'thinking': thinking,
            })
    return qa

edgecorner_object_templates = [
    "Please indicate the {name} located at {position} on the desktop.",
    "Please specify the {name} found at {position} on the desktop.",
    "Could you point out the {name} situated at {position} on the desktop?",
    "Kindly mark the {name} positioned at {position} on the desktop.",
    "Please identify the {name} located at the {position} on the desktop.",
    "Would you please indicate the {name} found at {position} on the desktop?",
    "Please locate the {name} that is at {position} on the desktop.",
    "Indicate the {name} that can be seen at {position} on the desktop.",
    "Could you please specify the {name} located at {position} on the desktop?",
    "Please denote the {name} positioned at the {position} area on the desktop.",
    "Point out the {name} that is positioned at {position} on the desktop, please.",
    "Would you kindly indicate the {name} that is situated at {position} on the desktop?",
    "Identify the {name} present at {position} on the desktop.",
    "Mark the {name} located at {position} on the desktop.",
    "Please reveal the {name} positioned at {position} on the desktop.",
    "Specify the {name} situated at {position} on the desktop.",
    "State the {name} found at {position} on the desktop.",
    "Please point out the {name} located in the {position} section of the desktop.",
    "Can you indicate which {name} is at {position} on the desktop?",
    "Kindly specify the {name} that is located at {position} on the desktop."
]
def add_info(dict, info):
    if 'object' in dict:
        dict['object'] = {}
    else:
        dict['object'] = info
    cate = info['cate']
    if cate in dict:
        dict[cate] = {}
    else:
        dict[cate] = info
def gen_edgecorner_object_qa(obj_dict, cam, table):
    qa = []
    dict = {
        'the far edge': {},
        'the near edge': {},
        'the left edge': {},
        'the right edge': {},
        'the far left corner': {},
        'the near left corner': {},
        'the far right corner': {},
        'the near right corner': {},
    }
    tabminx, tabminy, tabmaxx, tabmaxy, tablez = table
    for id, info in obj_dict.items():
        if not info['visible'] or not info['bottom_visible']:
            continue
        if info['bboxmin'][0] - tabminx < BOUND and info['bboxmin'][1] - tabminy < BOUND - 0.08:
            add_info(dict['the far left corner'], info)
        if tabmaxx - info['bboxmax'][0] < BOUND - 0.16 and info['bboxmin'][1] - tabminy < BOUND - 0.08:
            add_info(dict['the near left corner'], info)
        if info['bboxmin'][0] - tabminx < BOUND and tabmaxy - info['bboxmax'][1] < BOUND - 0.08:
            add_info(dict['the far right corner'], info)
        if tabmaxx - info['bboxmax'][0] < BOUND - 0.16 and tabmaxy - info['bboxmax'][1] < BOUND - 0.08:
            add_info(dict['the near right corner'], info)
        if info['bboxmin'][0] - tabminx < BOUND:
            add_info(dict['the far edge'], info)
        if tabmaxx - info['bboxmax'][0] < BOUND - 0.16:
            add_info(dict['the near edge'], info)
        if info['bboxmin'][1] - tabminy < BOUND - 0.08:
            add_info(dict['the left edge'], info)
        if tabmaxy - info['bboxmax'][1] < BOUND - 0.08:
            add_info(dict['the right edge'], info)
    for pos_name, dic in dict.items():
        for cate, info in dic.items():
            if info == {}:
                continue
            name = 'object' if cate == 'object' else info['minimal_desc']
            question = choice(edgecorner_object_templates).format(name=name, position=pos_name)
            if random() < 0.4:
                question = question + SUFFIX
            answer = f'[{format_coordinate(info)}]'
            thinking = ''
            qa.append({
                'question': question,
                'answer': answer,
                'thinking': thinking,
            })
            if cate == 'object':
                break
    return qa

edgecorner_point_templates = [
    "Please select any pixel point located at {position} of the desktop.",
    "Please choose any pixel point that lies at {position} on the desktop.",
    "Kindly pick a pixel point situated at {position} on the desktop.",
    "Select any pixel point that is positioned at {position} on the desktop.",
    "Would you please choose a pixel point located at {position} on the desktop?",
    "Please choose a pixel point found at {position} on the desktop.",
    "Identify any pixel point that is placed at {position} on the desktop.",
    "Pick any pixel point present at {position} on the desktop.",
    "Please mark a pixel point that is located at the {position} area of the desktop.",
    "Choose any pixel point situated at the {position} section of the desktop.",
    "Would you kindly select a pixel point at {position} on the desktop?",
    "Please indicate any pixel point positioned at {position} on the desktop.",
    "Select a pixel point located at {position} on the desktop, please.",
    "Please opt for any pixel point that is found at {position} on the desktop.",
    "Pick a pixel point that lies at {position} on the desktop.",
    "Kindly select a pixel point that is positioned at {position} of the desktop.",
    "Please choose any pixel point that resides at {position} on the desktop.",
    "Would you please pick a pixel point situated at {position} on the desktop?",
    "Select any pixel point that can be found at {position} on the desktop.",
    "Please choose a pixel point from {position} on the desktop."
]
def gen_edgecorner_point_qa(obj_dict, cam, table):
    qa = []
    dict = {
        'the far edge': [],
        'the near edge': [],
        'the left edge': [],
        'the right edge': [],
        'the far left corner': [],
        'the near left corner': [],
        'the far right corner': [],
        'the near right corner': [],
    }
    tabminx, tabminy, tabmaxx, tabmaxy, tablez = table
    for _ in range(30):
        rand_x = random() * (table[2] - table[0]) + table[0]
        rand_y = random() * (table[3] - table[1]) + table[1]
        if rand_x - tabminx < BOUND and rand_y - tabminy < BOUND - 0.08:
            dict['the far left corner'].append((rand_x, rand_y))
        elif tabmaxx - rand_x < BOUND - 0.16 and rand_y - tabminy < BOUND - 0.08:
            dict['the near left corner'].append((rand_x, rand_y))
        elif rand_x - tabminx < BOUND and tabmaxy - rand_y < BOUND - 0.08:
            dict['the far right corner'].append((rand_x, rand_y))
        elif tabmaxx - rand_x < BOUND - 0.16 and tabmaxy - rand_y < BOUND - 0.08:
            dict['the near right corner'].append((rand_x, rand_y))
        elif rand_x - tabminx < BOUND:
            dict['the far edge'].append((rand_x, rand_y))
        elif tabmaxx - rand_x < BOUND - 0.16:
            dict['the near edge'].append((rand_x, rand_y))
        elif rand_y - tabminy < BOUND - 0.08:
            dict['the left edge'].append((rand_x, rand_y))
        elif tabmaxy - rand_y < BOUND - 0.08:
            dict['the right edge'].append((rand_x, rand_y))
    for pos_name, points in dict.items():
        if points == []:
            continue
        for point in points:
            pix = check_table_pos(point, cam, table)
            if pix == None or pix_on_mask(pix, obj_dict):
                continue
            question = choice(edgecorner_point_templates).format(position=pos_name)
            if random() < 0.4:
                question = question + SUFFIX
            answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
            thinking = ''
            qa.append({
                'question': question,
                'answer': answer,
                'thinking': thinking,
            })
            break
    return qa

between_templates = [
    "Please select any pixel point on the desktop that lies between {object1} and {object2}.",
    "Please choose a pixel point on the desktop located between {object1} and {object2}.",
    "Kindly pick a pixel point on the desktop that is positioned between {object1} and {object2}.",
    "Could you mark any pixel point on the desktop that falls between {object1} and {object2}?",
    "Please indicate a pixel point on the desktop lying between {object1} and {object2}.",
    "Identify any pixel point on the desktop that is situated between {object1} and {object2}.",
    "Select a pixel point on the desktop that is located midway between {object1} and {object2}.",
    "Would you please choose a pixel point on the desktop that is found between {object1} and {object2}?",
    "Please point out a pixel point on the desktop that lies in the space between {object1} and {object2}.",
    "Can you indicate a pixel point on the desktop that falls in between {object1} and {object2}?",
    "Please pick any pixel point on the desktop that is set between {object1} and {object2}.",
    "Select any pixel coordinate on the desktop that is positioned between {object1} and {object2}.",
    "Mark a pixel point on the desktop that sits between {object1} and {object2}.",
    "Would you kindly mark a pixel point on the desktop that is in between {object1} and {object2}?",
    "Please designate a pixel point on the desktop that is located in the gap between {object1} and {object2}.",
    "Could you select a pixel point on the desktop that is nestled between {object1} and {object2}?",
    "Identify the pixel point on the desktop that lies in the middle between {object1} and {object2}.",
    "Please indicate the pixel coordinate on the desktop situated between {object1} and {object2}.",
    "Point out a pixel location on the desktop that is positioned in between {object1} and {object2}.",
    "Can you please choose a pixel point on the desktop located in the interval between {object1} and {object2}?"
]

def gen_between_qa(obj_dict, cam, table):
    qa = []
    if len(obj_dict) < 3:
        return qa
    times = 0
    while times < 10:
        times += 1
        keys = sample(list(obj_dict.keys()), k=2)
        obj1 = obj_dict[keys[0]]
        obj2 = obj_dict[keys[1]]
        if obj1['visible'] == False or obj2['visible'] == False or obj1['bottom_visible'] == False or obj2['bottom_visible'] == False:
            continue
        ref1 = choice(obj1['exclusive_ref'])
        name1 = ref_to_name(ref1)
        desc1 = ref_to_desc(ref1)
        ref2 = choice(obj2['exclusive_ref'])
        name2 = ref_to_name(ref2)
        desc2 = ref_to_desc(ref2)
        loc2d1 = obj1['2dloc']
        loc2d2 = obj2['2dloc']
        tabloc2d1 = obj1['2dtabloc']
        tabloc2d2 = obj2['2dtabloc']
        mid2d = ((tabloc2d1[0] + tabloc2d2[0]) / 2, (tabloc2d1[1] + tabloc2d2[1]) / 2)
        if pix_on_mask(mid2d, obj_dict):
            continue
        question = choice(between_templates).format(object1=desc1, object2=desc2)
        if random() < 0.4:
            question = question + SUFFIX
        thinking = f'''
{ref_to_thinking(obj_dict, obj1, ref1, desc1)}
{ref_to_thinking(obj_dict, obj2, ref2, desc2)}
{name1} at {format_coordinatexy(loc2d1[0], loc2d1[1])} is located at {format_coordinatexy(tabloc2d1[0], tabloc2d1[1])} on the table.
{name2} at {format_coordinatexy(loc2d2[0], loc2d2[1])} is located at {format_coordinatexy(tabloc2d2[0], tabloc2d2[1])} on the table.
We select one point within the interval as the answer: {format_coordinatexy(mid2d[0], mid2d[1])}.
        '''
        answer = f'[{format_coordinatexy(mid2d[0], mid2d[1])}]'
        qa.append({
            'question': question,
            'answer': answer,
            'thinking': thinking,
        })
    return qa

dist_rank_templates = [
    "Which {cate} is the {rank} closest to {object}?",
    "What is the {rank} closest {cate} to {object}?",
    "Which {cate} object ranks as the {rank} closest to {object}?",
    "Identify the {cate} that is the {rank} closest to {object}.",
    "Which {cate} ranks {rank} in proximity to {object}?",
    "Which {cate} is positioned as the {rank} nearest to {object}?",
    "Name the {cate} that is the {rank} closest to {object}.",
    "What {cate} object is the {rank} nearest to {object}?",
    "Find the {cate} that is the {rank} closest to {object}.",
    "Which {cate} can be found as the {rank} closest to {object}?",
    "Which {cate} is located as the {rank} nearest to {object}?",
    "Determine which {cate} stands as the {rank} closest to {object}.",
    "Which {cate} is the {rank} nearest neighbor to {object}?",
    "Select the {cate} that is the {rank} closest one to {object}.",
    "Identify the {rank} closest {cate} relative to {object}.",
    "Find out which {cate} is the {rank} closest to {object}.",
    "Which {cate} appears as the {rank} closest to {object}?",
    "What is the {cate} that is the {rank} nearest to {object}?",
    "Can you point out the {cate} that is the {rank} closest to {object}?",
    "Reveal the {cate} that stands as the {rank} closest to {object}?"
]
def cate_dist_rank(obj_dict, cate, color):
    sel_dict = {}
    for id, info in obj_dict.items():
        if not info['visible']:
            continue
        if cate != '' and info['cate'] != cate:
            continue
        if color != '' and info['color'] != color:
            continue
        sel_dict[id] = info
    while True:
        anchor = choice(list(obj_dict.keys()))
        if obj_dict[anchor]['visible']:
            break
    anchor_loc = obj_dict[anchor]['2dloc']
    ref = choice(obj_dict[anchor]['exclusive_ref'])
    desc = ref_to_desc(ref)
    count = len(sel_dict)
    if cate == '' or (obj_dict[anchor]['cate'] == cate and (color == '' or obj_dict[anchor]['color'] == color)):
        count -= 1
    rank = randint(1, count)
    question = choice(dist_rank_templates).format(cate=f'{color} {cate}', rank=get_rank(rank), object=desc)
    if random() < 0.4:
        question = question + SUFFIX
    id_dist = []
    thinking = ref_to_thinking(obj_dict, obj_dict[anchor], ref, desc) + '\n'
    for id, info in sel_dict.items():
        if id == anchor:
            continue
        # calc distance
        dist = calc_distance(obj_dict[anchor], info)
        id_dist.append((id, dist))
        thinking += f'{color} {cate} at {format_coordinate(info)} is {dist:3f}m away from it.\n'
    for id, dist in id_dist:
        closer_count = 0
        for id2, dist2 in id_dist:
            if id == id2:
                continue
            if dist2 < dist:
                closer_count += 1
        if closer_count == rank - 1:
            # found the rank
            answer = f'[{format_coordinate(obj_dict[id])}]'
            thinking += f'In summary, the {get_rank(rank)} closest {color} {cate} to {desc} is the object at {format_coordinate(obj_dict[id])}.\n'
            break
    return {
        'question': question,
        'answer': answer,
        'thinking': thinking,
    }
    
def gen_dist_rank_qa(obj_dict, cam, table, cates_count, color_cates_count):
    qa = []
    for _ in range(2):
        cate_dist_rank(obj_dict, '', '')
    for key, value in cates_count.items():
        if value < 3:
            continue
        for _ in range(2):
            cate_dist_rank(obj_dict, key, '')
    for key, value in color_cates_count.items():
        if value < 3:
            continue
        for _ in range(2):
            cate_dist_rank(obj_dict, key[0], key[1])
    return qa

two_constraints_templates = [
    "Please provide a pixel point on the desktop that simultaneously satisfies the following two spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Identify a pixel on the desktop that meets both of the following spatial criteria: 1. {condition1}; 2. {condition2}.",
    "Locate a pixel point on the desktop that fulfills these two spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Find a pixel on the desktop that adheres to both of the following spatial requirements: 1. {condition1}; 2. {condition2}.",
    "Please specify a pixel location on the desktop that satisfies the following two spatial constraints: 1. {condition1}; 2. {condition2}.",
    "Determine a pixel point on the desktop that meets these two spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Provide the coordinates of a pixel on the desktop that fulfills both of the following spatial criteria: 1. {condition1}; 2. {condition2}.",
    "Identify a pixel on the desktop that simultaneously meets these two spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Please locate a pixel point on the desktop that satisfies both of the following spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Find a pixel on the desktop that meets the following spatial requirements: 1. {condition1}; 2. {condition2}.",
    "Specify a pixel location on the desktop that adheres to both of these spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Determine a pixel on the desktop that fulfills the following two spatial criteria: 1. {condition1}; 2. {condition2}.",
    "Please identify a pixel point on the desktop that meets both of these spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Locate a pixel on the desktop that satisfies the following spatial constraints: 1. {condition1}; 2. {condition2}.",
    "Find a pixel point on the desktop that adheres to both of the following spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Provide a pixel location on the desktop that meets these two spatial criteria: 1. {condition1}; 2. {condition2}.",
    "Identify a pixel on the desktop that fulfills both of the following spatial requirements: 1. {condition1}; 2. {condition2}.",
    "Please specify a pixel point on the desktop that satisfies these two spatial conditions: 1. {condition1}; 2. {condition2}.",
    "Determine a pixel location on the desktop that meets both of the following spatial constraints: 1. {condition1}; 2. {condition2}.",
    "Locate a pixel on the desktop that adheres to the following two spatial conditions: 1. {condition1}; 2. {condition2}."
]
def get_dir(obj):
    dirs = [
        ('left', (0, -1)),
        ('right', (0, 1)),
        ('back', (-1, 0)),
        ('front', (1, 0)),
    ]
    dir = obj['direction']
    dir[0] = round(dir[0])
    dir[1] = round(dir[1])
    if not float_eq(abs(dir[0] + dir[1]), 1):
        dirid = randint(0, 3)
        return f'on the {dirs[dirid][0]} side of OBJECT', dirs[dirid][1]
    else:
        return obj['dir_desc'], dir

def gen_2constraints_qa(obj_dict, cam, table):
    qa = []
    for _ in range(20):
        obj1, obj2 = sample(list(obj_dict.keys()), k=2)
        ref1 = choice(obj_dict[obj1]['exclusive_ref'])
        desc1 = ref_to_desc(ref1)
        ref2 = choice(obj_dict[obj2]['exclusive_ref'])
        desc2 = ref_to_desc(ref2)
        thinking = f'''
{ref_to_thinking(obj_dict, obj_dict[obj1], ref1, desc1)}
Its position on the desktop is {format_coordinatexy(obj_dict[obj1]['2dtabloc'][0], obj_dict[obj1]['2dtabloc'][1])}.
{ref_to_thinking(obj_dict, obj_dict[obj2], ref2, desc2)}
Its position on the desktop is {format_coordinatexy(obj_dict[obj2]['2dtabloc'][0], obj_dict[obj2]['2dtabloc'][1])}.
'''
        dir_desc1, dir1 = get_dir(obj_dict[obj1])
        dir_desc2, dir2 = get_dir(obj_dict[obj2])
        if abs(dir1[0]) == abs(dir2[0]):
            continue
        if abs(dir1[0]) == 0:
            dir1, dir2 = dir2, dir1
            dir_desc1, dir_desc2 = dir_desc2, dir_desc1
            obj1, obj2 = obj2, obj1
        obj1 = obj_dict[obj1]
        obj2 = obj_dict[obj2]
        if dir1[0] > 0:
            if obj1['3dloc'][0] > obj2['3dloc'][0]:
                continue
        else:
            if obj1['3dloc'][0] < obj2['3dloc'][0]:
                continue
        if dir2[0] > 0:
            if obj1['3dloc'][1] < obj2['3dloc'][1]:
                continue
        else:
            if obj1['3dloc'][1] > obj2['3dloc'][1]:
                continue
        x = obj2['3dloc'][0]
        y = obj1['3dloc'][1]
        pix = world_to_camera_pixel((x, y, obj1['3dtabloc'][2]), cam)
        if (pix_on_mask((x, y), obj_dict)):
            continue
        question = choice(two_constraints_templates).format(
            condition1 = f'{dir_desc1.replace("OBJECT", desc1)}',
            condition2 = f'{dir_desc2.replace("OBJECT", desc2)}'
        )
        if random() < 0.4:
            question = question + SUFFIX
        answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
        thinking += 'The points in the corresponding directions of these two object is {format_coordinatexy(pix[0], pix[1])}.'
        qa.append({
            'question': question,
            'answer': answer,
            'thinking': thinking,
        })
    return qa


def gen_anchor_qa(obj_dict, cam, table, cates_count, color_cates_count):
    qa = []
    qa.extend(gen_orientation_distance_qa(obj_dict, cam, table))
    qa.extend(gen_edgecorner_object_qa(obj_dict, cam, table))
    qa.extend(gen_edgecorner_point_qa(obj_dict, cam, table))
    qa.extend(gen_between_qa(obj_dict, cam, table))
    qa.extend(gen_dist_rank_qa(obj_dict, cam, table, cates_count, color_cates_count))
    qa.extend(gen_2constraints_qa(obj_dict, cam, table))
    return qa
